import asyncio
import os
import json
from idlelib import query
from json import tool
from typing import Optional, List
from contextlib import AsyncExitStack
from datetime import datetime
import re

from openai import AsyncOpenAI
from dotenv import load_dotenv
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
load_dotenv()
class MCPClient:
    #最常用的初始化入口
    def __init__(self):
        #用于管理异步资源的生命周期（如关闭连接、释放锁等），确保安全退出；
        self.exit_stack = AsyncExitStack()
        self.openai_api_key = os.getenv('DATASCOPE_API_KEY')
        self.base_url = os.getenv('BASE_URL')
        self.model = os.getenv('MODEL')

        if not self.openai_api_key:
            raise ValueError("❌ 未找到 OpenAI API Key，请在 .env 文件中设置 DATASCOPE_API_KEY")

        #创建一个 OpenAI 兼容的 API 客户端实例
        # 并将其绑定到当前对象（self）上
        # 以便后续方法可以调用它来与大模型通信
        self.client =  AsyncOpenAI(api_key=self.openai_api_key,base_url=self.base_url)

        #self.session实例属性，用于管理与 MCP 工具服务器的连接会话
        #Optional[ClientSession]类型注解（Type Hint），表示该属性的值要么是 ClientSession 类型，要么是 None
        #（等价 于 Union[ClientSession, None]）
        self.session: Optional[ClientSession] = None


    async def connect_to_server(self,server_script_path: str):
        #对服务器脚本进行判断，只允许视.py 或者是.js文件
        is_python = server_script_path.endswith('.py')
        is_js = server_script_path.endswith('.js')

        if not (is_python or is_js):
            raise ValueError("服务器脚本必须是 .py 或 .js 文件")

        #确定启动命令， .py 用python，.js用node
        command = "python" if is_python else "node"

        #构造 MCP 所需要的服务器参数，包含启动命令，脚本路径参数，环境变量（为none）
        server_params = StdioServerParameters(command=command,args=[server_script_path],env=None)

        #启动 MCP 工具服务进程（并建立 stdio 通信）
        #使用 stdio_client 启动子进程，并将其注册到 exit_stack 中，确保退出时自动关闭。
        #返回值 stdio_transport 是一个元组 (reader, writer)，用于读写子进程的标准输入/输出。
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))

        #拆包通信通道，读取服务端返回的数据，并向服务端发送请求
        #self.stdio：用于从子进程读取数据（StreamReader）；
        #self.write：用于向子进程写入数据（StreamWriter）。
        self.stdio, self.write = stdio_transport

        #创建 MCP 客户端会话对象,并注册到资源栈中。
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio,self.write))

        #初始化会话,发送初始化消息给工具服务器，建立协议握手。
        await self.session.initialize()

        #获取工具列表并打印,（如 read_file, analyze_sentiment
        response = await self.session.list_tools()
        tools = response.tools
        print("\n已连接到服务器，支持以下工具:", [tool.name for tool in tools])

    #处理用户查询并执行工具链
    #接收用户自然语言查询，规划工具调用，执行，生成最终回复。
    async def process_query(self,query: str) ->str:

        #准备初始消息和获取工具列表,初始化对话历史，并再次获取工具列表（确保最新）
        messages = [{"role":"user","content":query}]
        response = await self.session.list_tools()

        #将 MCP 工具转换为 OpenAI 工具格式（虽然此处未启用 tool_choice="auto"，但保留结构）。
        available_tools = [
            {
                "type":"function",
                "function":{
                    "name": tool.name,
                    'description': tool.description,
                    'input_schema': tool.inputSchema
                }
            } for tool in response.tools
        ]

        #智能生成 Markdown 报告文件名：
        #提取问题的关键词，对文件名进行生成
        #在接收到用户提问后就应该生成出最后的输出的 md 文档的文件名
        #因为导出时若再生成文件名会导致部分组件无法识别该名称
        #从查询中提取关键词（如“关于特斯拉股价” → “特斯拉股价”）；re.search(pattern, string)：在字符串 query 中搜索第一个匹配正则表达式 pattern 的位置。
        keyword_match = re.search(r'(关于|分析|查询|搜索|查询)([^的\s，。？\n]+)',query)
        #清理非法文件名字符
        #如果re.search找到了匹配项（keyword_match不为None），就取第2个捕获组（即关键词）。
        #否则，使用默认关键词"分析对象"。这是一种安全的空值处理方式，避免程序崩溃
        keyword = keyword_match.group(2) if keyword_match else "分析对象"
        #移除 Windows/Linux/macOS 文件系统中非法的文件名字符。截取前 20 个字符，防止文件名过长
        self_keyword = re.sub(r'[\\/:*?"<>|]','',keyword)[:20]
        #添加时间戳避免冲突
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        #拼接完整路径（如 ./sentiment_reports/sentiment_特斯拉股价_20251109_150000.md）
        md_filename = f"sentiment_{self_keyword}_{timestamp}.md"
        #os.path.join()：跨平台安全地拼接路径（Windows 用 \，Linux/macOS 用 /）
        md_path = os.path.join("./sentiment_reports",md_filename)
        query = query.strip()+f"[md_filename={md_filename}][md_path={md_path}]"
        messages = [{"role": "user", "content": query}]

        #调用方法（通常是另一个 LLM 调用或规则引擎）,输出：一个工具调用计划，即一个步骤列表（list of dict）
        tool_plan = await self.plan_tool_usage(query,available_tools)
        print(f"🔧 工具调用计划: {tool_plan}")
        # 🔥 关键修复：确保 tool_plan 是 list
        if not isinstance(tool_plan, list):
            print(f"❌ 工具规划结果不是列表，而是 {type(tool_plan)}: {tool_plan}")
            return "❌ 任务规划失败，请重试。"

        tool_outputs = {} #用于存储每个已执行工具的输出结果
        messages = [{"role": "user", "content": query}] #再次初始化消息列表（冗余）


        tool_calls = []  # 存储模拟的 tool_calls
        tool_responses = []  # 存储对应的 tool 响应


        # 依次执行工具调用，并收集结果
        # [
        #     {"name": "analyze_sentiment", "arguments": {"text": "苹果公司最近表现很好"}},
        #     {"name": "send_email_with_attachment",
        #      "arguments": {"to": "user@example.com", "attachment_path": "{{analyze_sentiment}}"}}
        # ] 某个参数值形如 "{{analyze_sentiment}}"，说明它依赖前一个工具的输出
        for i,step in enumerate(tool_plan):
            if not isinstance(step, dict):
                print(f"⚠️ 跳过无效步骤（非字典）: {step}")
                continue
            tool_name = step["name"]
            tool_args = step["arguments"]
            if not tool_name or not isinstance(tool_args, dict):
                print(f"⚠️ 跳过无效工具调用: {step}")
                continue
            for key, val in tool_args.items():
                if isinstance(val, str) and val.startswith("{{") and val.endswith("}}"):
                    ref_key = val.strip("{} ")#去掉前后的大括号和空格
                    #从 tool_outputs 中查找该工具的输出结果：
                    #如果找到，就替换参数值；如果没找到（比如引用错误），保留原字符串（val）
                    resolved_val = tool_outputs.get(ref_key, val)
                    tool_args[key] = resolved_val
            # 注入统一的文件名或路径（用于分析和邮件）
            #硬编码的上下文注入逻辑，用于确保
            #1)情感分析工具知道要把结果写入哪个 Markdown 文件
            #2)发邮件工具知道附件路径是哪个
            #3)条件判断避免覆盖用户显式指定的参数
            if tool_name == "analyze_sentiment" and "filename" not in tool_args:
                tool_args["filename"] = md_filename
            if tool_name == "send_email_with_attachment" and "attachment_path" not in tool_args:
                tool_args["attachment_path"] = md_path
            #异步调用工具
            #result 通常是一个结构化响应对象，包含 .content 字段
            result = await self.session.call_tool(tool_name, tool_args)
            #缓存工具输出并记录到消息历史
            # output_text = result.content[0].text
            # 兼容两种情况：content 是字符串 或 是 [TextContent(...)]
            if isinstance(result.content, list):
                output_text = result.content[0].text if result.content else ""
            else:
                output_text = str(result.content)
            tool_outputs[tool_name] = output_text

            # 生成唯一 tool_call_id
            tool_call_id = f"call_{i}_{tool_name.replace('-', '_')}"
            # 记录 tool_call（用于 assistant 消息）
            tool_calls.append({
                "id": tool_call_id,
                "type": "function",
                "function": {
                    "name": tool_name,
                    "arguments": json.dumps(tool_args, ensure_ascii=False)
                }
            })
            # 记录 tool 响应
            tool_responses.append({
                "role": "tool",
                "tool_call_id": tool_call_id,  # 必须匹配！
                "content": output_text
            })

            # 构造合规的 messages
            messages = [{"role": "user", "content": query}]
            if tool_calls:
                # 添加 assistant 的 tool_calls 声明
                messages.append({
                    "role": "assistant",
                    "content": None,  # 有 tool_calls 时 content 应为 null
                    "tool_calls": tool_calls
                })
                # 添加所有 tool 返回结果
                messages.extend(tool_responses)
            # 否则直接让模型回答（无工具调用）
        # 调用大模型生成回复信息，并输出保存结果
        final_response = await  self.client.chat.completions.create(
            model=self.model,
            messages=messages
        )
        #从 API 响应中提取第一个候选回复的文本内容
        final_output = final_response.choices[0].message.content

        # 对辅助函数进行定义，目的是把文本清理成合法的文件名
        #内部工具函数，用于将任意字符串转换为操作系统兼容的文件名
        def clean_filename(text: str) -> str:
            text = text.strip()
            text = re.sub(r'[\\/:*?\"<>|]', '', text)
            return text[:50]

        # 使用清理函数处理用户查询，生成用于文件命名的前缀，并添加时间戳、设置输出目录
        # 最后构建出完整的文件路径用于保存记录
        safe_filename = clean_filename(query)
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"{safe_filename}_{timestamp}.txt"
        output_dir = "./llm_outputs"
        os.makedirs(output_dir, exist_ok=True)
        file_path = os.path.join(output_dir, filename)

        # 将对话内容写入 md 文档，其中包含用户的原始提问以及模型的最终回复结果
        #1)with open(...) as f：使用上下文管理器安全打开文件，确保写完后自动关闭。
        #2)"w" 模式：覆盖写入（每次都是新文件，不会追加）。
        #3)encoding="utf-8"：支持中文、表情符号等 Unicode 字符，避免乱码。
        #4)写入内容
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(f"🗣 用户提问：{query}\n\n")
            f.write(f"🤖 模型回复：\n{final_output}\n")

        print(f"📄 对话记录已保存为：{file_path}")

        return final_output

    #异步聊天主循环（chat loop），用于持续接收用户输入、调用 AI 处理逻辑，并输出回复。
    async def chat_loop(self):
        # 初始化提示信息
        print("\n🤖 MCP 客户端已启动！输入 'quit' 退出")

        # 进入主循环中等待用户输入
        while True:
            try:
                query = input("\n你: ").strip()#去除用户输入首尾的空白字符
                if query.lower() == 'quit':
                    break

                # 处理用户的提问，并返回结果
                response = await self.process_query(query)
                print(f"\n🤖 AI: {response}")

            except Exception as e:
                print(f"\n⚠️ 发生错误: {str(e)}")

    #大语言模型（LLM）根据用户自然语言请求，规划出一个有序的工具调用链
    async def plan_tool_usage(self, query: str, tools: List[dict]) -> List[dict]:
        # 构造系统提示词 system_prompt。
        # 将所有可用工具组织为文本列表插入提示中，并明确指出工具名，
        # 限定返回格式是 JSON，防止其输出错误格式的数据。
        print("\n📤 提交给大模型的工具定义:")
        #开发/调试阶段的日志，帮助开发者确认传给模型的工具是否正确
        #格式化 JSON 形式打印到控制台，ensure_ascii=False：允许输出中文等非 ASCII 字符，indent=2：美化缩进，便于阅读
        print(json.dumps(tools, ensure_ascii=False, indent=2))

        #1.
        #构建工具列表文本（用于系统提示），提取每个工具的：名称：tool['function']['name']描述：tool['function']['description']
        #拼接成多行字符串，eg.analyze_sentiment: 对一段文本进行情感分析，判断情绪倾向。
        tool_list_text = "\n".join([
            f"- {tool['function']['name']}: {tool['function']['description']}"
            for tool in tools
        ])

        #2.
        # 构造系统提示
        #构造一条 system role 消息
        # 用于指导模型行为。角色设定、输入、约束（防止幻觉）、输出格式强制要求
        system_prompt = {
            "role": "system",
            "content": (
                "你是一个智能任务规划助手，用户会给出一句自然语言请求。\n"
                "你只能从以下工具中选择（严格使用工具名称）：\n"
                f"{tool_list_text}\n"
                "如果多个工具需要串联，后续步骤中可以使用 {{上一步工具名}} 占位。\n"
                "返回格式：JSON 数组，每个对象包含 name 和 arguments 字段。\n"
                "不要返回自然语言，不要使用未列出的工具名。"
            )
        }

        #3.
        # 构造对话上下文并调用模型。
        # 将系统提示和用户的自然语言一起作为消息输入，并选用当前的模型。
        planning_messages = [
            system_prompt,
            {"role": "user", "content": query}
        ]

        #4.
        #调用大模型进行规划
        #tools=tools：传递工具定义（即使我们不希望模型自动调用工具，而是自己规划）。
        #tool_choice="none"：关键参数！  表示禁止模型自动触发工具调用 而是让模型在 content 中以自然语言或 JSON 形式输出计划
        #这是因为我们要的是“规划文本”，而不是直接执行工具
        response =  await  self.client.chat.completions.create(
            model=self.model,
            messages=planning_messages,
            tools=tools,
            tool_choice="none"
        )

        #5.
        # 尝试从 Markdown 代码块中提取 JSON
        # 很多模型喜欢用 Markdown 代码块 包裹 JSON 输出，
        # 例如 [{"name": "analyze_sentiment", "arguments": {...}}]
        content = response.choices[0].message.content.strip()
        match = re.search(r"```(?:json)?\\s*([\s\S]+?)\\s*```", content)
        if match:
            json_text = match.group(1) #group(1) 就是代码块内的 JSON 字符串
        else:
            json_text = content

        # 在解析 JSON 之后返回调用计划
        try:
            plan = json.loads(json_text) #尝试将字符串解析为 Python 对象
            return plan if isinstance(plan, list) else [] #解析结果是 list，说明是合法的步骤数组，直接返回
        except Exception as e:
            print(f"❌ 工具调用链规划失败: {e}\n原始返回: {content}")
            return []

    #这是一种优雅关闭的标准做法，避免资源泄漏
    async def cleanup(self):
        await self.exit_stack.aclose()


#一个顶层异步函数，作为程序的异步入口点
#所有异步操作（如连接服务器、聊天循环）都在这里协调执行
async def main():
    server_script_path = "D:\\桌面\\写的项目\\mcp-project\\server.py"
    client = MCPClient() #创建对象
    try:
        await client.connect_to_server(server_script_path)
        await client.chat_loop()
    finally:
        await client.cleanup()


if __name__ == "__main__":
    asyncio.run(main())



